import json
import time
import logging
import uuid
import cloudscraper
from typing import Dict, Any, AsyncGenerator, List

from fastapi import HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse

from app.core.config import settings
from app.providers.base_provider import BaseProvider
from app.utils.sse_utils import create_sse_data, create_chat_completion_chunk, DONE_CHUNK

logger = logging.getLogger(__name__)

class BritannicaProvider(BaseProvider):
    def __init__(self):
        self.scraper = cloudscraper.create_scraper()
        self.base_url = "https://www.britannica.com/chat-api"
        logger.info(f"BritannicaProvider 已初始化 ({settings.APP_VERSION})。")

    def _initialize_session(self) -> str:
        session_check_url = f"{self.base_url}/session/check"
        headers = {"Accept": "*/*", "Referer": "https://www.britannica.com/chatbot", "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36"}
        try:
            response = self.scraper.get(session_check_url, headers=headers, timeout=20)
            response.raise_for_status()
            data = response.json()
            session_id = data.get("session_id")
            if not session_id: raise ValueError("'/session/check' 响应中缺少 'session_id' 字段。")
            logger.info(f"成功获取一次性会话 ID: {session_id}")
            return session_id
        except Exception as e:
            raise RuntimeError(f"调用 /session/check 失败: {e}")

    def _create_conversation(self, session_id: str) -> str:
        create_conv_url = f"{self.base_url}/conversations/conversation/"
        headers = self._prepare_headers(session_id)
        payload = {"article_id": "-1"}
        try:
            response = self.scraper.post(create_conv_url, headers=headers, json=payload, timeout=20)
            response.raise_for_status()
            data = response.json()
            conv_id = data.get("id")
            if not conv_id: raise ValueError("'/conversations/conversation/' 响应中缺少 'id' 字段。")
            logger.info(f"成功创建新对话: {conv_id}")
            return conv_id
        except Exception as e:
            raise RuntimeError(f"创建新对话失败: {e}")

    def _initiate_chat(self, request_data: Dict[str, Any], session_id: str, conversation_id: str) -> Dict[str, str]:
        init_url = f"{self.base_url}/conversations/{conversation_id}/messages/"
        headers = self._prepare_headers(session_id)
        payload = self._prepare_payload(request_data)
        try:
            response = self.scraper.post(init_url, headers=headers, json=payload, timeout=settings.API_REQUEST_TIMEOUT)
            response.raise_for_status()
            data = response.json()
            run_id = data.get("run_id")
            if not run_id: raise ValueError("'/messages/' 响应中缺少 'run_id'。")
            logger.info(f"成功发送消息到对话 {conversation_id}, run_id: {run_id}")
            return {"run_id": run_id}
        except Exception as e:
            logger.error(f"发送消息失败: {e}", exc_info=True)
            raise HTTPException(status_code=500, detail=f"无法向上游发送消息: {str(e)}")

    async def chat_completion(self, request: Request):
        request_data = await request.json()
        use_stream = request_data.get("stream", False)
        
        try:
            session_id = self._initialize_session()
            conversation_id = self._create_conversation(session_id)
            init_data = self._initiate_chat(request_data, session_id, conversation_id)
            run_id = init_data.get("run_id")
        except Exception as e:
            logger.error(f"协议握手阶段失败: {e}", exc_info=True)
            raise HTTPException(status_code=503, detail=f"无法与上游服务建立会话: {e}")

        feed_url = f"{self.base_url}/conversations/{conversation_id}/messages/{run_id}/feed"
        headers = self._prepare_headers(session_id)

        if use_stream:
            return StreamingResponse(self._stream_response(feed_url, headers), media_type="text/event-stream")
        else:
            return await self._full_response(feed_url, headers)

    async def _stream_response(self, url: str, headers: Dict[str, str]) -> AsyncGenerator[bytes, None]:
        request_id = f"chatcmpl-{uuid.uuid4()}"
        try:
            response = self.scraper.get(url, headers=headers, stream=True, timeout=settings.API_REQUEST_TIMEOUT)
            response.raise_for_status()
            for line in response.iter_lines():
                if line.startswith(b"data:"):
                    content = line[len(b"data:"):].strip()
                    if content in (b"$$END$$", b"$$START$$"): continue
                    try:
                        data = json.loads(content)
                        if not isinstance(data, dict): continue
                        if data.get("name") in ["casual_response", "grounded_response"] and data.get("message_type") == "text_message_chunk":
                            delta_content = data.get("content")
                            if delta_content is not None:
                                yield create_sse_data(create_chat_completion_chunk(request_id, settings.DEFAULT_MODEL, delta_content))
                    except json.JSONDecodeError:
                        logger.warning(f"无法解析 SSE 数据块: {content}")
            yield create_sse_data(create_chat_completion_chunk(request_id, settings.DEFAULT_MODEL, "", "stop"))
            yield DONE_CHUNK
        except Exception as e:
            logger.error(f"处理流时发生错误: {e}", exc_info=True)
            error_chunk = create_chat_completion_chunk(request_id, settings.DEFAULT_MODEL, f"内部错误: {e}", "stop")
            yield create_sse_data(error_chunk)
            yield DONE_CHUNK

    async def _full_response(self, url: str, headers: Dict[str, str]) -> JSONResponse:
        request_id = f"chatcmpl-{uuid.uuid4()}"
        full_content_parts = []
        final_text = ""
        references = []
        try:
            response = self.scraper.get(url, headers=headers, stream=True, timeout=settings.API_REQUEST_TIMEOUT)
            response.raise_for_status()
            for line in response.iter_lines():
                if line.startswith(b"data:"):
                    content = line[len(b"data:"):].strip()
                    if content in (b"$$END$$", b"$$START$$"): continue
                    try:
                        data = json.loads(content)
                        if not isinstance(data, dict): continue
                        
                        if data.get("name") == "grounding_sources":
                            references = data.get("content", {}).get("search_results", [])

                        elif data.get("message_type") == "text_message":
                             text_from_final_message = data.get("content")
                             if isinstance(text_from_final_message, str):
                                final_text = text_from_final_message
                                # Do not break, continue to collect all data, just in case.
                        
                        elif data.get("message_type") == "text_message_chunk":
                            delta_content = data.get("content")
                            if delta_content:
                                full_content_parts.append(delta_content)
                    except json.JSONDecodeError:
                        continue
            
            if not final_text:
                final_text = "".join(full_content_parts)

            # Append references to the final text for non-streaming mode
            if references:
                final_text += "\n\n**References:**\n"
                for i, ref in enumerate(references):
                    final_text += f"[{i+1}] {ref.get('title', 'N/A')}: <{ref.get('url', '#')}>\n"

            response_data = {"id": request_id, "object": "chat.completion", "created": int(time.time()), "model": settings.DEFAULT_MODEL, "choices": [{"index": 0, "message": {"role": "assistant", "content": final_text}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 0, "completion_tokens": len(final_text), "total_tokens": len(final_text)}}
            return JSONResponse(response_data)
        except Exception as e:
            logger.error(f"处理聚合响应时发生错误: {e}", exc_info=True)
            raise HTTPException(status_code=500, detail=f"内部服务器错误: {e}")

    def _prepare_headers(self, session_id: str) -> Dict[str, str]:
        return {"Accept": "application/json, text/plain, */*", "Accept-Language": "en-US,en;q=0.9", "Content-Type": "application/json", "Origin": "https://www.britannica.com", "Referer": "https://www.britannica.com/chatbot", "session-id": session_id, "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36"}

    def _prepare_payload(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
        messages = request_data.get("messages", [])
        last_user_message = next((m['content'] for m in reversed(messages) if m['role'] == 'user'), "Hello")
        return {"article_id": "-1", "role": "user", "content": {"message_type": "text_message", "text_message": last_user_message}, "action": "run", "retry_status": "standard"}

    async def get_models(self) -> JSONResponse:
        model_data = {"object": "list", "data": [{"id": name, "object": "model", "created": int(time.time()), "owned_by": "lzA6"} for name in settings.KNOWN_MODELS]}
        return JSONResponse(content=model_data)
